# newspaper_scraper.py
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException, NoSuchElementException, WebDriverException
from bs4 import BeautifulSoup # 确保这里是 from bs4 import BeautifulSoup
import time
import os
import json
import pandas as pd
import argparse
import platform
import traceback # 用于打印更详细的错误堆栈

# 全局变量，用于存储ChromeDriver路径
CHROME_DRIVER_PATH = None

def set_chromedriver_path(path):
    global CHROME_DRIVER_PATH
    CHROME_DRIVER_PATH = path
    print(f"INFO: ChromeDriver path set to: {CHROME_DRIVER_PATH}")

def get_default_chromedriver_path():
    """
    尝试获取相对于当前脚本的 ChromeDriver 路径。
    假设 ChromeDriver 可执行文件位于脚本所在目录下的 'chromedriver-win64' 子目录中。
    """
    try:
        # 获取当前脚本 (newspaper_scraper.py) 所在的绝对目录
        script_dir = os.path.dirname(os.path.abspath(__file__))
    except NameError:
        # 如果 __file__ 未定义 (例如在某些交互式环境或打包后)，尝试使用当前工作目录
        # 但这可能不是期望的脚本目录，所以最好在脚本正常执行时依赖 __file__
        script_dir = os.getcwd()
        print(f"WARN: __file__ not defined, using current working directory for relative path: {script_dir}")


    relative_subdir = "chromedriver-win64" # ChromeDriver 所在的子目录名

    if platform.system() == "Windows":
        chromedriver_filename = "chromedriver.exe"
    elif platform.system() == "Darwin": # macOS
        chromedriver_filename = "chromedriver"
    else: # Linux 等
        chromedriver_filename = "chromedriver"
        
    # 构建相对于脚本目录的完整路径
    potential_path = os.path.join(script_dir, relative_subdir, chromedriver_filename)

    # 检查这个相对路径是否存在且文件可执行
    if os.path.isfile(potential_path): # isfile 检查是否是文件
        if platform.system() != "Windows": # Windows 不需要显式检查X权限，isfile足够
            if os.access(potential_path, os.X_OK):
                print(f"INFO: Found default ChromeDriver via relative path (executable): {potential_path}")
                return potential_path
            else:
                print(f"WARN: Default ChromeDriver found at relative path but NOT executable: {potential_path}. Check permissions.")
                return None # 文件存在但不可执行
        else: # Windows
            print(f"INFO: Found default ChromeDriver via relative path: {potential_path}")
            return potential_path
    else:
        print(f"WARN: Default ChromeDriver not found at relative path: {potential_path}. Will rely on PATH or UI config.")
        return None

def get_newspaper_details(newspaper_name, driver_path_override=None):
    search_url = "https://www.nppa.gov.cn/bsfw/cyjghcpcx/bzbs/index.html"
    driver = None
    # 初始化details字典，确保所有期望的键都存在，即使是查询失败
    details = {
        "查询报纸": newspaper_name,
        "机构名称": "未获取", # 默认值
        "主管单位": "未获取", # 默认值
        "主办单位": "未获取"  # 默认值
    }

    current_chromedriver_path = driver_path_override if driver_path_override else CHROME_DRIVER_PATH
    if not current_chromedriver_path:
        current_chromedriver_path = get_default_chromedriver_path()

    try:
        options = webdriver.ChromeOptions()
        options.add_argument("--headless")
        options.add_argument("--disable-gpu")
        options.add_argument("--no-sandbox")
        options.add_argument('--log-level=3')
        options.add_experimental_option('excludeSwitches', ['enable-logging'])
        # options.add_argument("--window-size=1920,1080")

        if current_chromedriver_path and os.path.exists(current_chromedriver_path):
            print(f"INFO: Using ChromeDriver at: {current_chromedriver_path}")
            service = webdriver.chrome.service.Service(executable_path=current_chromedriver_path)
            driver = webdriver.Chrome(service=service, options=options)
        elif current_chromedriver_path:
             print(f"ERROR: ChromeDriver not found at specified path: {current_chromedriver_path}. Will attempt to use PATH.")
             try:
                 driver = webdriver.Chrome(options=options)
                 print("INFO: Successfully started ChromeDriver from PATH.")
             except WebDriverException as e_path:
                 print(f"ERROR: ChromeDriver not found at {current_chromedriver_path} AND failed to start from PATH: {e_path}")
                 details.update({"机构名称": "错误", "主管单位": "ChromeDriver未找到或启动失败", "主办单位": ""})
                 return details
        else:
            print("INFO: No ChromeDriver path specified, attempting to find in PATH.")
            try:
                driver = webdriver.Chrome(options=options)
                print("INFO: Successfully started ChromeDriver from PATH.")
            except WebDriverException as e_path:
                print(f"ERROR: Failed to start ChromeDriver from PATH: {e_path}")
                details.update({"机构名称": "错误", "主管单位": "ChromeDriver未能从PATH启动", "主办单位": ""})
                return details
        
        if not driver:
            details.update({"机构名称": "错误", "主管单位": "ChromeDriver未能初始化", "主办单位": ""})
            return details

        driver.set_page_load_timeout(45)
        driver.set_script_timeout(30)

        print(f"INFO: Navigating to search URL for '{newspaper_name}'...")
        driver.get(search_url)
        
        print(f"INFO: Locating name input field for '{newspaper_name}'...")
        name_input_field = WebDriverWait(driver, 20).until(
            EC.presence_of_element_located((By.ID, "jgmc"))
        )
        name_input_field.send_keys(newspaper_name)
        print(f"INFO: Entered '{newspaper_name}' into name field.")
        
        print(f"INFO: Locating captcha field for '{newspaper_name}'...")
        captcha_value_element = WebDriverWait(driver, 20).until(
            EC.presence_of_element_located((By.ID, "checkCode"))
        )
        captcha_text = captcha_value_element.get_attribute("value")
        print(f"INFO: Captcha value found: '{captcha_text}' for '{newspaper_name}'.")
        
        captcha_input_field = driver.find_element(By.ID, "input1")
        captcha_input_field.send_keys(captcha_text)
        print(f"INFO: Entered captcha for '{newspaper_name}'.")
        
        search_button = driver.find_element(By.ID, "query")
        search_button.click()
        print(f"INFO: Clicked search button for '{newspaper_name}'.")
        
        try:
            print(f"INFO: Waiting for '查看详细' link for '{newspaper_name}'...")
            view_details_link = WebDriverWait(driver, 20).until(
                EC.element_to_be_clickable((By.LINK_TEXT, "查看详细"))
            )
            view_details_link.click()
            print(f"INFO: Clicked '查看详细' for '{newspaper_name}'.")
        except TimeoutException:
            error_message = "未找到'查看详细'链接或超时"
            try:
                no_results_message_xpath = "//td[contains(text(), '无符合条件的查询结果')]"
                WebDriverWait(driver, 5).until( 
                    EC.presence_of_element_located((By.XPATH, no_results_message_xpath))
                )
                error_message = "无符合条件的查询结果"
                print(f"INFO: No results found for '{newspaper_name}'.")
            except (NoSuchElementException, TimeoutException):
                print(f"ERROR: Timeout waiting for '查看详细' or '无结果' message for '{newspaper_name}'.")
            details.update({"机构名称": "查询失败", "主管单位": error_message, "主办单位": ""})
            return details

        print(f"INFO: Waiting for details page to load for '{newspaper_name}'...")
        WebDriverWait(driver, 20).until(
            EC.presence_of_element_located((By.XPATH, "//td[contains(text(), '机构名称')]"))
        )
        time.sleep(0.5)

        print(f"INFO: Extracting details for '{newspaper_name}'...")
        page_source = driver.page_source
        soup = BeautifulSoup(page_source, 'html.parser')

        rows = soup.find_all('tr')
        data_to_extract = {"机构名称": "未获取", "主管单位": "未获取", "主办单位": "未获取"} # 重置以防万一
        for row in rows:
            cells = row.find_all('td')
            if len(cells) == 2:
                label = cells[0].get_text(strip=True)
                value = cells[1].get_text(strip=True)
                if label in data_to_extract: # 只提取我们关心的字段
                    data_to_extract[label] = value if value else "空" # 如果值为空字符串，标记为"空"
        
        details.update(data_to_extract)
        print(f"INFO: Successfully extracted details for '{newspaper_name}': {data_to_extract}")
        return details

    except WebDriverException as e:
        error_msg = f"WebDriver错误: {str(e)[:200]}" # 稍微加长错误信息
        details.update({"机构名称": "错误", "主管单位": error_msg, "主办单位": ""})
        print(f"ERROR: WebDriverException for '{newspaper_name}': {e}")
        return details
    except TimeoutException:
        details.update({"机构名称": "错误", "主管单位": "页面/元素加载超时", "主办单位": ""})
        print(f"ERROR: Page/element loading timed out for '{newspaper_name}'.")
        return details
    except NoSuchElementException:
        details.update({"机构名称": "错误", "主管单位": "页面元素未找到", "主办单位": ""})
        print(f"ERROR: Could not find a specified page element for '{newspaper_name}'.")
        return details
    except Exception as e:
        error_msg = f"未知错误: {str(e)[:200]}"
        details.update({"机构名称": "错误", "主管单位": error_msg, "主办单位": ""})
        print(f"ERROR: An unexpected error occurred for '{newspaper_name}': {e}")
        traceback.print_exc() 
        return details
    finally:
        if driver:
            print(f"INFO: Quitting driver for '{newspaper_name}'...")
            driver.quit()
            print(f"INFO: Driver for '{newspaper_name}' quit.")

def save_results_to_excel(results_list, filename="newspaper_query_results.xlsx", output_dir=None):
    if not results_list:
        print("INFO: 没有结果可以保存到 Excel。")
        return

    columns = ["查询报纸", "机构名称", "主管单位", "主办单位"]
    processed_results = []
    for res_dict in results_list:
        # 确保每个结果字典都包含所有列，即使是失败的查询
        # get_newspaper_details 在失败时也会返回包含这些键的字典
        record = {col: res_dict.get(col, "未获取/错误") for col in columns}
        processed_results.append(record)

    df = pd.DataFrame(processed_results, columns=columns)
    
    filepath = filename # 默认保存到当前目录
    if output_dir:
        if not os.path.exists(output_dir):
            try:
                os.makedirs(output_dir)
                print(f"INFO: 创建输出目录: {output_dir}")
            except OSError as e:
                print(f"ERROR: 创建输出目录 '{output_dir}' 失败: {e}. 将保存到当前目录。")
                # output_dir = None # 保持 filepath 默认值
        
        # 只有当 output_dir 成功创建或已存在时，才拼接路径
        if os.path.exists(output_dir): # 再次检查
            base_filename, ext = os.path.splitext(filename)
            safe_base_filename = "".join(c if c.isalnum() or c in (' ', '_', '-') else '_' for c in base_filename).rstrip()
            safe_filename = f"{safe_base_filename}{ext}"
            filepath = os.path.join(output_dir, safe_filename)
    
    try:
        df.to_excel(filepath, index=False, engine='openpyxl')
        print(f"SUCCESS: 查询结果已成功保存到文件: {os.path.abspath(filepath)}")
    except Exception as e:
        print(f"ERROR: 保存 Excel 文件失败: {e}")
        traceback.print_exc()

def process_single_query_non_interactive(target_newspaper, output_dir=None, driver_path=None):
    print(f"INFO: 开始单个查询 (非交互模式): {target_newspaper}")
    if not target_newspaper.strip():
        print("ERROR: 报纸名称不能为空。")
        # 对于单个查询，如果名称为空，可能不应该生成空的Excel，所以这里可以return
        # 或者，也可以生成一个只有查询报纸列为空的Excel行
        # 为了与批量查询一致，我们还是获取一个空的details并保存
        details = {"查询报纸": "", "机构名称": "错误", "主管单位": "报纸名称为空", "主办单位": ""}
        results_list = [details]
        # 后续会调用 save_results_to_excel
    else:
        if driver_path:
            set_chromedriver_path(driver_path)
        newspaper_info = get_newspaper_details(target_newspaper.strip())
        results_list = [newspaper_info]

    print(f"\n--- 查询结果 for {target_newspaper if target_newspaper.strip() else '空名称'} ---")
    # 确保 newspaper_info 是上面 if/else 分支中的一个
    current_info_to_print = results_list[0] 
    print(f"  查询报纸: {current_info_to_print.get('查询报纸', target_newspaper.strip())}")
    print(f"  机构名称: {current_info_to_print.get('机构名称', '未找到')}")
    print(f"  主管单位: {current_info_to_print.get('主管单位', '未找到')}")
    print(f"  主办单位: {current_info_to_print.get('主办单位', '未找到')}")
    
    if results_list: 
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        # 使用更安全的名称处理，即使 target_newspaper 为空
        safe_newspaper_name = "".join(c if c.isalnum() else '_' for c in target_newspaper.strip() or "empty_name")
        filename = f"single_query_{safe_newspaper_name}_{timestamp}.xlsx"
        save_results_to_excel(results_list, filename=filename, output_dir=output_dir)
        
        if current_info_to_print.get("机构名称") in ["错误", "查询失败"] or not target_newspaper.strip() :
            print(f"INFO: '{target_newspaper if target_newspaper.strip() else '空名称'}' 查询未成功或名称为空，但结果已记录到Excel。")
        else:
            print(f"INFO: '{target_newspaper}' 查询成功，结果已保存到Excel。")
    else:
        print(f"INFO: 未对 '{target_newspaper if target_newspaper.strip() else '空名称'}' 进行任何查询操作或获取任何结果，未保存Excel文件。")

def process_batch_query(json_file_path="piliang.json", output_dir=None, driver_path=None):
    try:
        with open(json_file_path, 'r', encoding='utf-8') as f:
            raw_data = json.load(f)
    except FileNotFoundError:
        print(f"ERROR：找不到 JSON 文件 '{json_file_path}'。")
        return
    except json.JSONDecodeError:
        print(f"ERROR：无法解析 JSON 文件 '{json_file_path}'。")
        return

    if driver_path:
        set_chromedriver_path(driver_path)

    all_results_for_excel = []
    newspaper_names_to_query = []

    if isinstance(raw_data, list):
        for item in raw_data:
            if isinstance(item, str) and item.strip():
                newspaper_names_to_query.append(item.strip())
            else:
                print(f"WARN：JSON 文件中的列表项 '{item}' 无效或为空，已跳过。")
    elif isinstance(raw_data, dict):
        print("INFO: 检测到JSON文件为字典格式，尝试按 name_newspaper_X 键读取。")
        for i in range(1, 500): # 稍微增加上限
            key = f"name_newspaper_{i}"
            if key in raw_data:
                name = raw_data[key]
                if isinstance(name, str) and name.strip():
                    newspaper_names_to_query.append(name.strip())
                else:
                    print(f"WARN：JSON 文件中键 '{key}' 的值无效或为空，已跳过。")
            elif i > 1 and f"name_newspaper_{i-1}" not in raw_data:
                break
    else:
        print(f"ERROR: JSON文件 '{json_file_path}' 的格式无法识别。")
        return
            
    if not newspaper_names_to_query:
        print("INFO: 在 JSON 文件中没有找到有效的报纸名称进行查询。")
        # 如果JSON为空或无效，但文件存在，也可以选择生成一个空的Excel或只包含表头的Excel
        # 当前行为是不生成Excel，因为 all_results_for_excel 会为空
        return

    print(f"\n--- 开始批量查询（共 {len(newspaper_names_to_query)} 个报纸）---")
    
    for index, newspaper_name in enumerate(newspaper_names_to_query):
        print(f"\n[{index+1}/{len(newspaper_names_to_query)}] 正在查询: {newspaper_name}")
        if not newspaper_name.strip(): # 处理JSON中可能存在的空字符串条目
            print(f"WARN: 跳过查询，因为报纸名称为空。")
            details_dict = {
                "查询报纸": "", 
                "机构名称": "错误", 
                "主管单位": "报纸名称为空", 
                "主办单位": ""
            }
        else:
            details_dict = get_newspaper_details(newspaper_name) 
        
        print(f"  查询报纸: {details_dict.get('查询报纸', newspaper_name)}")
        print(f"  机构名称: {details_dict.get('机构名称', '未获取')}")
        print(f"  主管单位: {details_dict.get('主管单位', '未获取')}")
        print(f"  主办单位: {details_dict.get('主办单位', '未获取')}")

        all_results_for_excel.append(details_dict)
        if index < len(newspaper_names_to_query) - 1:
            print("INFO: 等待1秒后进行下一次查询...")
            time.sleep(1) 

    print("\n\n--- 批量查询完成 ---")
    
    if all_results_for_excel:
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        excel_filename = f"batch_query_results_{timestamp}.xlsx"
        save_results_to_excel(all_results_for_excel, filename=excel_filename, output_dir=output_dir)
        
        success_count = sum(1 for res in all_results_for_excel if res.get("机构名称") not in ["错误", "查询失败", "未获取"] and res.get("查询报纸", "").strip() != "")
        fail_count = len(all_results_for_excel) - success_count
        print(f"INFO: 批量查询中，有效成功 {success_count} 条，失败或无效条目 {fail_count} 条。所有结果已尝试记录到Excel。")
    else:
        print("INFO: 批量查询没有产生任何结果 (可能是JSON列表为空或所有条目均无效)。")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="报纸信息查询工具 (命令行)")
    parser.add_argument(
        "--mode", 
        choices=['single', 'batch', 'interactive'], 
        default='interactive', 
        help="查询模式: 'single' (单个查询), 'batch' (批量查询), 'interactive' (交互模式，默认)"
    )
    parser.add_argument("--name", type=str, help="单个查询模式下的报纸名称")
    parser.add_argument("--json_file", type=str, default="piliang.json", help="批量查询模式下使用的JSON文件路径 (默认: piliang.json)")
    parser.add_argument("--output_dir", type=str, help="Excel文件保存的目录 (可选)")
    parser.add_argument("--chromedriver_path", type=str, help="ChromeDriver的可执行文件路径 (可选)")

    args = parser.parse_args()

    if args.chromedriver_path:
        set_chromedriver_path(args.chromedriver_path)
    
    if args.mode == 'single':
        # process_single_query_non_interactive 会处理名称为空的情况
        process_single_query_non_interactive(args.name if args.name else "", args.output_dir, args.chromedriver_path)
    elif args.mode == 'batch':
        process_batch_query(args.json_file, args.output_dir, args.chromedriver_path)
    elif args.mode == 'interactive':
        print("INFO: 进入交互式查询模式...")
        print("--- 欢迎使用报纸查询工具 (交互模式) ---")
        while True:
            choice = input("请选择操作模式：\n1. 单个查询\n2. 批量查询 (从 piliang.json)\nQ. 退出\n请输入选项 (1/2/Q): ").strip().lower()
            if choice == '1':
                target_newspaper = input("请输入要查询的报纸名称: ")
                process_single_query_non_interactive(target_newspaper) 
            elif choice == '2':
                json_path_interactive = input(f"请输入JSON文件路径 (默认为 piliang.json): ") or "piliang.json"
                process_batch_query(json_path_interactive)
            elif choice == 'q':
                print("正在退出程序...")
                break
            else:
                print("无效的选项，请重新输入。")
    else:
        print(f"ERROR: 未知的模式 '{args.mode}'")
